#!/usr/bin/env python3
"""
批量处理蛋模型文件的脚本
为所有类型蛋添加CustomModelData支持
"""

import os
import json
from pathlib import Path

# 基础类型列表
BASE_TYPES = [
    "bug", "dark", "dragon", "electric", "fairy", "fighting", 
    "flying", "ghost", "grass", "ground", "ice", "normal", 
    "poison", "psychic", "rock", "steel", "water"
]

# 已处理的类型（避免重复处理）
PROCESSED_TYPES = {"pokemon", "manaphy", "fire", "water", "grass", "electric", "dragon", "steel"}

def get_overrides_json():
    """返回标准的overrides JSON结构"""
    return [
        {"predicate": {"custom_model_data": 1}, "model": "cobbreeding:item/{}_4v"},
        {"predicate": {"custom_model_data": 2}, "model": "cobbreeding:item/{}_5v"},
        {"predicate": {"custom_model_data": 3}, "model": "cobbreeding:item/{}_6v"},
        {"predicate": {"custom_model_data": 4}, "model": "cobbreeding:item/{}_shiny"},
        {"predicate": {"custom_model_data": 5}, "model": "cobbreeding:item/{}_shiny_4v"},
        {"predicate": {"custom_model_data": 6}, "model": "cobbreeding:item/{}_shiny_5v"},
        {"predicate": {"custom_model_data": 7}, "model": "cobbreeding:item/{}_shiny_6v"}
    ]

def create_variant_model(egg_type, variant, has_shiny=False):
    """创建变体模型文件"""
    if has_shiny:
        if variant == "shiny_4v":
            return {
                "parent": "item/generated",
                "textures": {
                    "layer0": f"cobbreeding:item/{egg_type}_pokemon_egg",
                    "layer1": "cobbreeding:item/overlay_4v",
                    "layer2": "cobbreeding:item/overlay_shiny"
                }
            }
        elif variant == "shiny_5v":
            return {
                "parent": "item/generated",
                "textures": {
                    "layer0": f"cobbreeding:item/{egg_type}_pokemon_egg",
                    "layer1": "cobbreeding:item/overlay_5v",
                    "layer2": "cobbreeding:item/overlay_shiny"
                }
            }
        elif variant == "shiny_6v":
            return {
                "parent": "item/generated",
                "textures": {
                    "layer0": f"cobbreeding:item/{egg_type}_pokemon_egg",
                    "layer1": "cobbreeding:item/overlay_6v",
                    "layer2": "cobbreeding:item/overlay_shiny"
                }
            }
        else:  # shiny
            return {
                "parent": "item/generated",
                "textures": {
                    "layer0": f"cobbreeding:item/{egg_type}_pokemon_egg",
                    "layer1": "cobbreeding:item/overlay_shiny"
                }
            }
    else:
        overlay = f"overlay_{variant}"
        return {
            "parent": "item/generated",
            "textures": {
                "layer0": f"cobbreeding:item/{egg_type}_pokemon_egg",
                "layer1": f"cobbreeding:item/{overlay}"
            }
        }

def process_egg_type(models_dir, egg_name):
    """处理单个蛋类型"""
    egg_file = models_dir / f"{egg_name}_pokemon_egg.json"
    
    if not egg_file.exists():
        print(f"跳过 {egg_name}: 文件不存在")
        return
    
    # 读取现有文件
    try:
        with open(egg_file, 'r', encoding='utf-8') as f:
            data = json.load(f)
    except Exception as e:
        print(f"错误读取 {egg_file}: {e}")
        return
    
    # 检查是否已有overrides
    if "overrides" in data:
        print(f"跳过 {egg_name}: 已有overrides")
        return
    
    # 添加overrides
    overrides = get_overrides_json()
    for override in overrides:
        override["model"] = override["model"].format(f"{egg_name}_pokemon_egg")
    
    data["overrides"] = overrides
    
    # 写回文件
    try:
        with open(egg_file, 'w', encoding='utf-8') as f:
            json.dump(data, f, indent=2, ensure_ascii=False)
        print(f"✅ 处理完成: {egg_name}")
    except Exception as e:
        print(f"错误写入 {egg_file}: {e}")
        return
    
    # 创建变体模型文件
    variants = ["4v", "5v", "6v", "shiny", "shiny_4v", "shiny_5v", "shiny_6v"]
    
    for variant in variants:
        variant_file = models_dir / f"{egg_name}_pokemon_egg_{variant}.json"
        
        if variant_file.exists():
            continue  # 跳过已存在的文件
        
        has_shiny = variant.startswith("shiny")
        clean_variant = variant.replace("shiny_", "")
        
        variant_data = create_variant_model(egg_name, variant, has_shiny)
        
        try:
            with open(variant_file, 'w', encoding='utf-8') as f:
                json.dump(variant_data, f, indent=2, ensure_ascii=False)
            print(f"  ➤ 创建变体: {variant}")
        except Exception as e:
            print(f"  ❌ 创建变体失败 {variant}: {e}")

def main():
    """主函数"""
    script_dir = Path(__file__).parent
    models_dir = script_dir / "common" / "src" / "main" / "resources" / "assets" / "cobbreeding" / "models" / "item"
    
    if not models_dir.exists():
        print(f"错误: 模型目录不存在: {models_dir}")
        return
    
    print("🚀 开始批量处理蛋模型文件...")
    print(f"📁 模型目录: {models_dir}")
    
    # 处理单类型蛋
    print("\n📦 处理单类型蛋:")
    for egg_type in BASE_TYPES:
        if egg_type in PROCESSED_TYPES:
            print(f"跳过 {egg_type}: 已手动处理")
            continue
        process_egg_type(models_dir, egg_type)
    
    # 处理双类型蛋
    print("\n📦 处理双类型蛋:")
    processed_dual = 0
    for type1 in BASE_TYPES:
        for type2 in BASE_TYPES:
            if type1 != type2:
                dual_name = f"{type1}_{type2}"
                process_egg_type(models_dir, dual_name)
                processed_dual += 1
                
                # 显示进度
                if processed_dual % 50 == 0:
                    print(f"  进度: 已处理 {processed_dual} 个双类型蛋...")
    
    print(f"\n🎉 批量处理完成!")
    print(f"📊 总计处理了 {len(BASE_TYPES) - len(PROCESSED_TYPES)} 个单类型蛋")
    print(f"📊 总计处理了 {processed_dual} 个双类型蛋")
    print("\n📋 接下来需要做的:")
    print("1. 制作4个叠加材质文件:")
    print("   - overlay_4v.png (绿色)")
    print("   - overlay_5v.png (紫色)")
    print("   - overlay_6v.png (青色)")
    print("   - overlay_shiny.png (金色)")
    print("2. 测试游戏中的效果")

if __name__ == "__main__":
    main() 